import os
import matplotlib.pyplot as plt
import numpy as np
import cv2
import pandas as pd

def fig2data(fig):
    fig.savefig('matplotlib.png')
    fig = cv2.imread('matplotlib.png')
    os.remove('matplotlib.png')
    return fig

def EqualizeHist(original):
    b, g, r = cv2.split(original)
    colors = ('b', 'g', 'r')
    
    b1 = cv2.equalizeHist(b)
    g1 = cv2.equalizeHist(g)
    r1 = cv2.equalizeHist(r)

    fig = plt.figure()
    plt.title("’Flattened’ Color Histogram")
    plt.xlabel("Bins")
    plt.ylabel("# of Pixels")

    HIST = []

    for (chan, color) in zip([b, g, r], colors):
        hist = cv2.calcHist([chan], [0], None, [256], [0, 256])
        plt.plot(hist, color = color)
        HIST.append(hist.flatten())
        plt.xlim([0, 256])
        
    df = pd.DataFrame(np.array(HIST), index=colors, columns=np.arange(0,256))

    return fig2data(fig), cv2.merge([b1,g1,r1]), cv2.merge([b,g,r1]), cv2.merge([b,g1,r]), cv2.merge([b1,g,r]), df

def ColorFilter(original, colorLow, colorHigh):
    hsv = cv2.cvtColor(original, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv, np.array(colorLow), np.array(colorHigh))
    kernal = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernal)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernal)
    reverse_mask = cv2.bitwise_not(mask)
    return cv2.bitwise_and(original, original, mask=mask), cv2.bitwise_and(original, original, mask=reverse_mask)
